package logger

import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"sync/atomic"
	"time"
)

// contextKey 自定义 context key 类型
type contextKey string

const (
	traceIDKey     contextKey = "traceID"
	transactionKey contextKey = "transactionNo"
)

var (
	traceCounter uint64
	logger       *slog.Logger
)

// 初始化日志
func init() {
	// 创建 JSON 格式的日志处理器
	handler := slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
		Level: slog.LevelDebug,
	})
	logger = slog.New(handler)
}

// 生成 traceID
func generateTraceID() string {
	counter := atomic.AddUint64(&traceCounter, 1)
	return time.Now().Format("20060102150405") + "-" + fmt.Sprintf("%06d", counter)
}

// 创建带 traceID 的上下文
func WithTraceID(ctx context.Context) context.Context {
	traceID := generateTraceID()
	return context.WithValue(ctx, traceIDKey, traceID)
}

// 从上下文获取 traceID
func GetTraceID(ctx context.Context) string {
	if traceID, ok := ctx.Value(traceIDKey).(string); ok {
		return traceID
	}
	return "unknown"
}

// 增加事务编号
func IncrementTransactionNo(ctx context.Context) context.Context {
	no := GetTransactionNo(ctx) + 1
	return context.WithValue(ctx, transactionKey, no)
}

// 获取事务编号
func GetTransactionNo(ctx context.Context) int {
	if no, ok := ctx.Value(transactionKey).(int); ok {
		return no
	}
	return 0
}

// 记录调试日志
func LogDebug(ctx context.Context, msg string, args ...any) {
	logger.Debug(msg, append([]any{
		"traceID", GetTraceID(ctx),
		"transactionNo", GetTransactionNo(ctx),
	}, args...)...)
}

// 记录信息日志
func LogInfo(ctx context.Context, msg string, args ...any) {
	logger.Info(msg, append([]any{
		"traceID", GetTraceID(ctx),
		"transactionNo", GetTransactionNo(ctx),
	}, args...)...)
}

// 记录警告日志
func LogWarn(ctx context.Context, msg string, args ...any) {
	logger.Warn(msg, append([]any{
		"traceID", GetTraceID(ctx),
		"transactionNo", GetTransactionNo(ctx),
	}, args...)...)
}

// 记录错误日志
func LogError(ctx context.Context, msg string, args ...any) {
	logger.Error(msg, append([]any{
		"traceID", GetTraceID(ctx),
		"transactionNo", GetTransactionNo(ctx),
	}, args...)...)
}

// 记录协议日志
func LogProtocol(ctx context.Context, msg string, args ...any) {
	logger.Info("[PROTOCOL] "+msg, append([]any{
		"traceID", GetTraceID(ctx),
		"transactionNo", GetTransactionNo(ctx),
	}, args...)...)
}
